LinkLuaModifier("modifier_cut_trees", "modifiers/modifier_cut_trees", LUA_MODIFIER_MOTION_NONE)

UnitEntity = UnitEntity or class({}, nil, Entity)

function UnitEntity:constructor(unitName, location, findClearSpace, owner, teamNumber)
	getbase(UnitEntity).constructor(self)

	if unitName ~= nil then
		if owner and owner:GetUnit() and IsValidEntity(owner:GetUnit()) then
			owner = owner:GetUnit()
		end
		self:SetUnit(CreateUnitByName(unitName, location, findClearSpace, owner, owner, teamNumber))
	end

	self.nettable = {}
	self.removeOnDeath = true
end

function UnitEntity:AddComponent(component)
    if component.Activate then
        component.Activate(self)
    end

    for key, value in pairs(component) do
        if key ~= "Activate" then
            if self[key] ~= nil then
                print("Warning! Colliding key", key)
            end

            self[key] = value
        end
    end

    return self
end

function UnitEntity:SetUnit(unit)
	if unit ~= nil and unit:IsNPC() then
		return
	end

	self.unit = unit

	self.unit.GetParentEntity = function(unit)
		return self
	end
end

function UnitEntity:IsAlive()
	return self:GetUnit():IsAlive()
end

function UnitEntity:GetUnit()
	return self.unit
end

function UnitEntity:GetName()
	return self:GetUnit():GetName()
end

function UnitEntity:IsValid()
	return self:GetUnit() ~= nil and IsValidEntity(self:GetUnit())
end

function UnitEntity:SetPos(pos)
	if not self:IsValid() then
		return
	end

	self:GetUnit():SetAbsOrigin(pos)
end

function UnitEntity:GetPos()
	if not self:IsValid() then
		return
	end

	return self:GetUnit():GetAbsOrigin()
end

function UnitEntity:GetTeamNumber()
	return self:GetUnit():GetTeamNumber()
end

function UnitEntity:GetFacing()
	if not self:IsValid() then
		return nil
	end

	return self:GetUnit():GetForwardVector()
end

function UnitEntity:SetFacing(facing)
	if not self:IsValid() then
		return
	end

	self:GetUnit():SetForwardVector(facing)
end

function UnitEntity:StartGesture(nActivity)
	if not self:IsValid() then
		return
	end

	self:GetUnit():StartGesture(nActivity)
end

function UnitEntity:StartGestureWithPlaybackRate(nActivity, flRate)
	if not self:IsValid() then
		return
	end

	self:GetUnit():StartGestureWithPlaybackRate(nActivity, flRate)
end

function UnitEntity:RemoveGesture(nActivity)
	if not self:IsValid() then
		return
	end

	self:GetUnit():RemoveGesture(nActivity)
end

function UnitEntity:EmitSound(sound, location)
	if not self:IsValid() then
		return
	end

	if location then
		EmitSoundOnLocationWithCaster(location, sound, self:GetUnit())
	else
		self:GetUnit():EmitSound(sound)
	end
end

function UnitEntity:StopSound(sound)
	if not self:IsValid() then
		return
	end

	self:GetUnit():StopSound(sound)
end

function UnitEntity:AddNewModifier(source, ability, modifierName, params)
	if not self:IsValid() then
		return
	end

	if source and source:GetUnit() and IsValidEntity(source:GetUnit()) then
		source = source:GetUnit()
	end

	return self:GetUnit():AddNewModifier(source, ability, modifierName, params)
end

function UnitEntity:HasModifier(modifierName)
	if not self:IsValid() then
		return false
	end

	return self:GetUnit():HasModifier(modifierName)
end

function UnitEntity:RemoveModifierByName(modifierName)
	if not self:IsValid() then
		return
	end

	self:GetUnit():RemoveModifierByName(modifierName)
end

function UnitEntity:RemoveModifierByNameAndCaster(modifierName, caster)
	if not self:IsValid() then
		return
	end
	if caster and caster:GetUnit() and IsValidEntity(caster:GetUnit()) then
		caster = caster:GetUnit()
	end

	self:GetUnit():RemoveModifierByNameAndCaster(modifierName, caster)
end

function UnitEntity:FindModifierByName(modifierName)
	if not self:IsValid() then
		return nil
	end

	return self:GetUnit():FindModifierByName(modifierName)
end

function UnitEntity:FindModifierByNameAndCaster(modifierName, caster)
	if not self:IsValid() then
		return nil
	end

	if caster and caster:GetUnit() and IsValidEntity(caster:GetUnit()) then
		caster = caster:GetUnit()
	end

	return self:GetUnit():FindModifierByNameAndCaster(modifierName, caster)
end

function UnitEntity:FindAllModifiers()
	if not self:IsValid() then
		return nil
	end

	return self:GetUnit():FindAllModifiers()
end

function UnitEntity:FindAllModifiersByName(modifierName)
	if not self:IsValid() then
		return nil
	end

	return self:GetUnit():FindAllModifiersByName(modifierName)
end

function UnitEntity:GetGroundHeight(position)
	if not self:IsValid() then
		return nil
	end

	return GetGroundHeight(position or self:GetPos(), self:GetUnit())
end

function UnitEntity:FindClearSpace(position, force)
	if not self:IsValid() then
		return
	end

	FindClearSpaceForUnit(self:GetUnit(), position, force)
end

function UnitEntity:GetMaxHealth()
	if not self:IsValid() then
		return
	end

	return self:GetUnit():GetMaxHealth()
end

function UnitEntity:Damage(source, amount, isPhysical)
end

function UnitEntity:Heal(amount, source)
	if not self:IsValid() then
		return
	end

	if amount == nil then
		return
	end

	self:GetUnit():Heal(amount, source)
end

function UnitEntity:DropItemAtPositionImmediate(item, position)
	if not self:IsValid() then
		return
	end

	if position == nil then
		position = self:GetPos()
	end

	self:GetUnit():DropItemAtPositionImmediate(item:GetItem(), position)
end

function UnitEntity:SetCanCutTrees(bAble)
    if bAble then
        self:AddNewModifier(self, nil, "modifier_cut_trees", {})
    else
        self:RemoveModifierByName("modifier_cut_trees")
    end
end

function UnitEntity:CanCutTrees()
    return self:HasModifier("modifier_cut_trees")
end

function UnitEntity:SetNetTable(key, value)
	self.nettable = self.nettable or {}

	self.nettable[key] = value
	self:UpdateNetTable()
end

function UnitEntity:GetNetTable(key)
	self.nettable = self.nettable or {}

	return self.nettable[key]
end

function UnitEntity:UpdateNetTable()
	self.nettable = self.nettable or {}

	CustomNetTables:SetTableValue("unitEntities", tostring(self:GetUnit():entindex()), self.nettable)
end

function UnitEntity:Remove()
	self.nettable = nil
	CustomNetTables:SetTableValue("unitEntities", tostring(self:GetUnit():entindex()), nil)

	if self:IsValid() then
		if self.removeOnDeath then
			self:GetUnit():RemoveSelf()
		end
	end

	getbase(UnitEntity).Remove(self)
end